package com.senke.aiagent.tools;

import cn.hutool.json.JSONUtil;
import com.alibaba.cloud.ai.toolcalling.baidusearch.BaiduSearchService;
import lombok.extern.slf4j.Slf4j;
import org.jsoup.Jsoup;
import org.jsoup.nodes.Document;
import org.jsoup.nodes.Element;
import org.jsoup.select.Elements;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.tool.annotation.Tool;
import org.springframework.ai.tool.annotation.ToolParam;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Mono;

import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;

/**
 * 百度搜索工具
 */
@Slf4j
public class BaiduSearchTool {

    // Baidu 的搜索接口地址
    private static final String SEARCH_API_URL = "https://www.baidu.com/s?wd=";

    private final WebClient webClient = WebClient.builder().defaultHeader("User-Agent", new String[]{"Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36"}).defaultHeader("Accept", new String[]{"text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7"}).defaultHeader("Accept-Encoding", new String[]{"gzip, deflate"}).defaultHeader("Content-Type", new String[]{"application/x-www-form-urlencoded"}).defaultHeader("Referer", new String[]{"https://www.baidu.com/"}).defaultHeader("Accept-Language", new String[]{"zh-CN,zh;q=0.9,ja;q=0.8"}).codecs((configurer) -> {
        configurer.defaultCodecs().maxInMemorySize(5242880);
    }).build();

    private static final Logger logger = LoggerFactory.getLogger(BaiduSearchTool.class);

    @Tool(description = "Search for information from Baidu Search Engine", returnDirect = false)
    public String searchWeb(@ToolParam(description = "Search query keyword", required = true) String query) {
        if (StringUtils.hasText(query)) {
            String url = SEARCH_API_URL + query;
            try {
                Mono<String> responseMono = this.webClient.get().uri(url, new Object[0]).retrieve().bodyToMono(String.class);
                String html = (String) responseMono.block();
                assert html != null;
                List<BaiduSearchService.SearchResult> results = this.parseHtml(html);
                if (CollectionUtils.isEmpty(results)) {
                    return null;
                } else {
                    logger.info("baidu search: {},result number:{}", query, results.size());
                    Iterator var7 = results.iterator();
                    while (var7.hasNext()) {
                        BaiduSearchService.SearchResult d = (BaiduSearchService.SearchResult) var7.next();
                        Logger var10000 = logger;
                        String var10001 = d.title();
                        var10000.info(var10001 + "\n" + d.abstractText());
                    }
                    List<BaiduSearchService.SearchResult> searchResults = results.subList(0, Math.min(results.size(), 5));
                    return JSONUtil.toJsonStr(searchResults);
                }
            } catch (Exception var9) {
                Exception e = var9;
                logger.error("failed to invoke baidu search caused by:{}", e.getMessage());
                return null;
            }
        } else {
            return null;
        }
    }

    private List<BaiduSearchService.SearchResult> parseHtml(String htmlContent) {
        try {
            Document doc = Jsoup.parse(htmlContent);
            Element contentLeft = doc.selectFirst("div#content_left");
            Elements divContents = contentLeft.children();
            List<BaiduSearchService.SearchResult> listData = new ArrayList();
            Iterator var6 = divContents.iterator();

            while (true) {
                String title;
                String abstractText;
                while (true) {
                    Element div;
                    do {
                        if (!var6.hasNext()) {
                            return listData;
                        }

                        div = (Element) var6.next();
                    } while (!div.hasClass("c-container"));

                    title = "";
                    abstractText = "";

                    try {
                        if (!div.hasClass("xpath-log") && !div.hasClass("result-op")) {
                            if (!"se_com_default".equals(div.attr("tpl"))) {
                                continue;
                            }

                            if (div.selectFirst("h3") != null) {
                                title = div.selectFirst("h3").text().trim();
                            } else {
                                title = ((Element) div.children().get(0)).text().trim();
                            }

                            if (div.selectFirst("div.c-abstract") != null) {
                                abstractText = div.selectFirst("div.c-abstract").text().trim();
                            } else if (div.selectFirst("div") != null) {
                                abstractText = div.selectFirst("div").text().trim();
                            } else {
                                abstractText = div.text().trim();
                            }
                            break;
                        }

                        if (div.selectFirst("h3") != null) {
                            title = div.selectFirst("h3").text().trim();
                        } else {
                            title = div.text().trim().split("\n", 2)[0];
                        }

                        if (div.selectFirst("div.c-abstract") != null) {
                            abstractText = div.selectFirst("div.c-abstract").text().trim();
                        } else if (div.selectFirst("div") != null) {
                            abstractText = div.selectFirst("div").text().trim();
                        } else {
                            abstractText = div.text().trim().split("\n", 2)[1].trim();
                        }
                        break;
                    } catch (Exception var11) {
                        Exception e = var11;
                        logger.error("failed to parse baidu search html,caused by:{}", e.getMessage());
                    }
                }

                listData.add(new BaiduSearchService.SearchResult(title, abstractText));
            }
        } catch (Exception var12) {
            Exception e = var12;
            logger.error("failed to parse baidu search html,caused by:{}", e.getMessage());
            return null;
        }
    }
}
